# _*_ coding: utf-8 _*_
'''
author:小龙ssss
date: 2025/02/06
'''
import wx  # pip install wxPython
import requests
from run_in_new_shell import run_in_new_shell


class ChatFrame(wx.Frame):
    """
    定义一个用于展示聊天界面的窗口类。

    构思思路：
    - 通过继承wx.Frame类创建自定义窗口，便于扩展和维护。
    - 使用wx.Panel作为容器来放置其他控件，保持界面整洁有序。
    - 利用布局管理器（如BoxSizer）灵活地管理控件的位置和大小。
    """

    def __init__(self, initial_model_name: str):
        """
        初始化ChatFrame实例。

        参数:
            initial_model_name (str): 初始使用的模型名称。
        """
        # 调用父类构造函数初始化窗口，并设置窗口标题为'奥拉玛聊天机器人'
        super().__init__(parent=None, title='Ollama 聊天机器人', size=(800, 600))

        # 主面板设置，所有控件都将放置在此面板上
        main_panel = wx.Panel(self)

        # 布局管理器，垂直布局用于整个界面，水平布局用于输入框和按钮
        vertical_layout = wx.BoxSizer(wx.VERTICAL)
        horizontal_layout = wx.BoxSizer(wx.HORIZONTAL)

        # 创建显示对话历史的文本控件，用户可以看到与AI的对话记录
        # 同时要确保TextCtrl支持富文本，即可以设置字体颜色、背景色等
        # 更新后的代码，其中我们确保了聊天历史记录的wx.TextCtrl控件启用了富文本支持（通过设置TE_RICH或TE_RICH2样式）
        self.conversation_history = wx.TextCtrl(
            main_panel,
            style=wx.TE_MULTILINE | wx.TE_READONLY | wx.HSCROLL | wx.TE_RICH2,
            size=(-1, 400)
        )

        # 用户输入框，用于输入要发送给AI的消息
        self.user_input_box = wx.TextCtrl(
            main_panel,
            style=wx.TE_MULTILINE,
            size=(-1, 100)
        )

        # 发送按钮，点击后将用户输入的内容发送给AI并显示回复
        self.send_button = wx.Button(main_panel, label='发送')
        self.send_button.Bind(wx.EVT_BUTTON, self.handle_send)

        # 设置布局，将输入框和发送按钮加入水平布局，然后将水平布局和对话历史加入垂直布局
        horizontal_layout.Add(self.user_input_box, proportion=1, flag=wx.EXPAND | wx.RIGHT, border=5)
        horizontal_layout.Add(self.send_button, flag=wx.EXPAND)

        vertical_layout.Add(self.conversation_history, proportion=1, flag=wx.EXPAND | wx.ALL, border=5)
        vertical_layout.Add(horizontal_layout, flag=wx.EXPAND | wx.LEFT | wx.RIGHT | wx.BOTTOM, border=5)

        main_panel.SetSizer(vertical_layout)

        # API配置，定义请求的URL和默认使用的模型名称
        self.api_URL = 'http://localhost:11434/api/generate'
        self.selected_model = initial_model_name

    def append_text_with_color(self, text_ctrl, text, color):
        """
        在指定的wx.TextCtrl控件中以特定颜色添加文本。

        参数:
            text_ctrl: 目标wx.TextCtrl控件。
            text: 要添加的文本内容。
            color: 文本颜色（如"blue", "green"等）。
        """
        # 创建一个临时的TextAttr对象来设置颜色
        text_attr = wx.TextAttr(wx.Colour(color))
        # 应用样式并追加文本
        text_ctrl.SetDefaultStyle(text_attr)
        text_ctrl.AppendText(text + "\n")
        # 重置样式
        text_ctrl.SetDefaultStyle(wx.TextAttr())


    def handle_send(self, event):
        """处理用户点击发送按钮的动作"""
        user_message = self.user_input_box.GetValue().strip()
        if not user_message:
            return

        # 清空输入框以便用户输入新消息
        self.user_input_box.SetValue('')
        # 将用户消息添加到对话历史中
        self.append_text_with_color(self.conversation_history, f'小龙ssss: {user_message}', "blue")
        # self.conversation_history.AppendText(f'小龙ssss: {user_message}\n')

        try:
            # prompt参数表示输入的文本，stream参数表示是否使用流式响应，这里设置为False
            request_payload = {'model': self.selected_model, 'prompt': user_message, 'stream': False}
            response = requests.post(self.api_URL, json=request_payload)

            if response.status_code == 200:
                ai_reply = response.json().get('response', '未收到回复')
                # 将AI的回复添加到对话历史中
                self.conversation_history.AppendText(f'AI: {ai_reply}\n\n')
            else:
                self.append_text_with_color(self.conversation_history, f'错误: 获取回复失败 (状态码: {response.status_code})',
                                            "red")
                # self.conversation_history.AppendText(f'错误: 获取回复失败 (状态码: {response.status_code})\n\n')
        except Exception as e:
            self.append_text_with_color(self.conversation_history, f'错误: {str(e)}', "red")
            # self.conversation_history.AppendText(f'错误: {str(e)}\n\n')

        # 自动滚动到底部，确保最新消息可见
        self.conversation_history.ShowPosition(self.conversation_history.GetLastPosition())


def launch_application():
    app = wx.App()
    # 这里使用的是qwen2.5的模型，可以改用其它已经下载好的模型
    chat_window = ChatFrame(initial_model_name='qwen2.5:latest')
    chat_window.Show()
    app.MainLoop()


if __name__ == '__main__':
    # 启动Ollama服务和聊天窗口
    run_in_new_shell('ollama_server_start.py')
    launch_application()
